import sys
parent_dir = "../pygl"
sys.path.append(parent_dir)
import OpenGL.GL as gl
import glm
import imgui
from glfw import _GLFWwindow as GLFWwindow
import numpy as np
import re
import pathlib
import pygl
from common import AppWithTrackball, AxesHelper, GridHelper
from volume_cube import VolumeCube

class VolumeViewer(AppWithTrackball):
    # 初始化着色器，vao vbo等
    def setupVertex(self) -> None:
        gl.glEnable(gl.GL_DEPTH_TEST)
        gl.glEnable(gl.GL_BLEND)
        gl.glBlendFunc(gl.GL_SRC_ALPHA, gl.GL_ONE_MINUS_SRC_ALPHA)
        self.program = pygl.ProgramVF("usecase/shaders/volume_raycast.vs", 
                                      "usecase/shaders/volume_raycast.fs")
        self.background_color = glm.vec4(0.5, 0.5, 1.0, 1.0)
        raw_file = "dataset/aneurism_256x256x256_uint8.raw"
        pattern = re.compile(r'(.+)_(\d+)x(\d+)x(\d+)_(.+).raw')
        raw_info = pattern.match(pathlib.Path(raw_file).name)
        raw_shape = int(raw_info.group(4)), int(raw_info.group(3)), int(raw_info.group(2))
        raw_dtype = raw_info.group(5)
        with open(raw_file, 'rb') as f:
            raw_data = np.fromfile(f,  dtype=raw_dtype).reshape(raw_shape)
        self.volume_cube = VolumeCube(raw_data)
        self.imgui_window_width = 250
        self.imgui_window_height = 200
        self.max_shape = np.max(raw_shape)
        self.step = 1.0/self.max_shape
        self.step_size = glm.vec3(self.step)
        self.threshold = 0.2
        self.camera.fov = 60
        self.camera.reset(glm.vec3(1.0, 1.0, 1.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        self.world_axes = AxesHelper(1.0, 5.0)
        self.grid = GridHelper(2.0, 20, glm.vec3(1.0, 0.0, 0.0), 
                               glm.vec3(0.0, 0.0, 1.0))
        self.show_world_axes = True
        self.show_grid = True

    # 配置imgui界面元素
    def showImGui(self)->None:
        imgui.begin("设置")
        if imgui.button("重置轨迹球相机"):
            self.camera.reset(glm.vec3(1.0, 1.0, 1.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        _, self.step = imgui.slider_float("步长大小", self.step, 0.001, 0.01)
        self.step_size = glm.vec3(self.step)
        _, self.threshold = imgui.slider_float("阈值大小", self.threshold, 0.01, 0.99)
        _, self.show_world_axes = imgui.checkbox("显示三维世界坐标轴", 
                                                self.show_world_axes)
        _, self.show_grid = imgui.checkbox("显示XZ平面网格", 
                                                self.show_grid)
        imgui.end()       

    # 渲染图形
    def render(self) -> None:
        gl.glClearBufferfv(gl.GL_COLOR, 0, self.background_color.to_tuple())
        gl.glClearBufferfv(gl.GL_DEPTH, 0, 1.0)
        projection = self.camera.getProjectionMatrix()
        view = self.camera.getViewMatrix()
        model = glm.mat4(1.0)
        mv = view * model
        camera_position = glm.vec3(glm.inverse(mv) * glm.vec4(0, 0, 0, 1))
        self.program.use()
        self.program.setUniformMatrix4fv(0, model)
        self.program.setUniformMatrix4fv(1, view)
        self.program.setUniformMatrix4fv(2, projection)
        self.program.setUniform1i(3, 0)
        self.program.setUniform3fv(4, camera_position)
        self.program.setUniform3fv(5, self.step_size)
        self.program.setUniform1f(6, self.threshold)
        self.volume_cube.render()
        gl.glDisable(gl.GL_BLEND)
        if self.show_world_axes:
            self.world_axes.updateModelMatrix
            self.world_axes.render(view, projection)
        if self.show_grid:
            self.grid.render(view, projection)

    # 释放资源
    def cleanup(self) -> None:
        self.program.delete()
        self.grid.delete()
        self.world_axes.delete()
        self.volume_cube.delete()

    def mouse_move(self, window:GLFWwindow, x:float, y:float)->None:
        if imgui.is_any_item_active():
            return
        super().mouse_move(window, x, y)
    
if __name__ == "__main__":
    app = VolumeViewer(100, 100, 1280, 920, "你好，体渲染App！")
    app.run()


